#include <iostream>
#include <winsock2.h>
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")

#define CLOSE_SOCKET closesocket

// ICMP 报文结构
struct icmp_header
{
  uint8_t type;
  uint8_t code;
  uint16_t checksum;
  uint16_t id;
  uint16_t seq;
};

// 计算校验和
uint16_t checksum(void *buf, int len)
{
  uint32_t sum = 0;
  auto *data = (uint16_t *)buf;
  while (len > 1)
  {
    sum += *data++;
    len -= 2;
  }
  if (len == 1)
  {
    sum += *(uint8_t *)data;
  }
  sum = (sum >> 16) + (sum & 0xFFFF);
  sum += (sum >> 16);
  return (uint16_t)~sum;
}

int main()
{
  // 初始化 Winsock
  WSADATA wsa;
  if (WSAStartup(MAKEWORD(2, 2), &wsa) != 0)
  {
    std::cerr << "WSAStartup failed: " << WSAGetLastError() << std::endl;
    return 1;
  }

  // 创建原始套接字（需要管理员权限）
  SOCKET sock = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
  if (sock == INVALID_SOCKET)
  {
    std::cerr << "Socket creation error: " << WSAGetLastError() << std::endl;
    WSACleanup();
    return 1;
  }

  // 绑定到所有接口
  sockaddr_in addr{};
  addr.sin_family = AF_INET;
  addr.sin_addr.s_addr = INADDR_ANY;
  if (bind(sock, (sockaddr *)&addr, sizeof(addr)) == SOCKET_ERROR)
  {
    std::cerr << "Bind failed: " << WSAGetLastError() << std::endl;
    CLOSE_SOCKET(sock);
    WSACleanup();
    return 1;
  }

  std::cout << "ICMP Server listening (Run as Administrator)..." << std::endl;

  while (true)
  {
    char buffer[1024];
    sockaddr_in from{};
    int fromlen = sizeof(from);

    // 接收ICMP报文
    int rcvd = recvfrom(sock, buffer, sizeof(buffer), 0,
                        (sockaddr *)&from, &fromlen);
    if (rcvd == SOCKET_ERROR)
    {
      std::cerr << "Receive error: " << WSAGetLastError() << std::endl;
      continue;
    }

    // 解析 ICMP 报文（跳过20字节IP头）
    if (rcvd < 20 + sizeof(icmp_header))
    {
      std::cerr << "Invalid packet size" << std::endl;
      continue;
    }

    icmp_header *icmp = (icmp_header *)(buffer + 20);

    // 只处理 Echo Request (Type=8)
    if (icmp->type == 8)
    {
      // Windows 下使用 inet_ntoa 替代 inet_ntop
      char *ipstr = inet_ntoa(from.sin_addr);
      std::cout << "Ping request from " << ipstr << std::endl;

      // 修改为 Echo Reply (Type=0)
      icmp->type = 0;

      // 重算校验和（必须置零后计算）
      icmp->checksum = 0;
      icmp->checksum = checksum(icmp, rcvd - 20);

      // 发送响应
      if (sendto(sock, (char *)icmp, rcvd - 20, 0,
                 (sockaddr *)&from, fromlen) == SOCKET_ERROR)
      {
        std::cerr << "Send failed: " << WSAGetLastError() << std::endl;
      }
    }
  }

  CLOSE_SOCKET(sock);
  WSACleanup();
  return 0;
}